from pprint import pprint
from pyspark import Row
from pyspark.ml import Estimator
from pyspark.ml.linalg import Vectors
from pyspark.ml.param.shared import HasInputCols, HasOutputCols, HasOutputCol
from pyspark.sql import DataFrame, Window, functions
from typing import List, Callable, Dict

from gai.v2.spark.base import SparkTransformer
from gai.v2.spark.feature.woe_iv_calculator import compute_group_woe_iv
from gai.v2.spark.feature.variable_stats import _compute_scores_for_categorical_variable_binary_class, \
    _compute_scores_for_categorical_variable_multiclass, explode_list_plain, explode_map_plain, prepare_test_dataframe


class FeatureScoreCalculator(SparkTransformer):
    """Computes information value (for binary class only),
    information gain and chi square stat for specified features.

    Args:
        labelCol: the name of label column
        labelRange: the collection of values that a label can take
        compoCols: the list of names of composite columns. A composite column represents a collection of features.
        compoExploders: the list of callables that disassemble each cell of composite columns into lists. The order of
            exploders should match that of ``compoCols``.
        simpleCols: the list of names of simple columns. A simple column represents a single feature.

    Examples:

    Consider a sample as follows::

        +----+------+------------+------------------+-----+
        | age|  name|    app_list|     category_list|label|
        +----+------+------------+------------------+-----+
        |12.0| Alice| app_1,app_3| cate_0:3,cate_1:2|    0|
        |13.0|   Bob|app_3,app_21|cate_0:4,cate_10:7|    1|
        |14.0|Claire| app_1,app_7|cate_0:3,cate_11:4|    0|
        +----+------+------------+------------------+-----+

    In this example, ``app_list`` and ``category_list`` are composite columns, while ``age`` is a simple column.
    The following code snippet instantiates and uses a ``FeatureScoreCalculator`` for ``app_list``, ``category_list``, and ``age``.
    For now we require each exploder to return a list of tuples each in the form of ``(feature_name, feature_value)``.
    Take the first row of the above dataframe for example, ``explode_list_plain`` is expected to return
    ``[('app_1',1.0),('app_3',1.0)]``, and ``explode_map_plain`` to return ``[('cate_0',3.0),('cate_1',2.0)]``.

    Notes:
        It is important for the exploder to return a list of pairs whose second elements
        are ``float``.

    >>> df = prepare_test_dataframe()
    >>> df.show()
    +---+------+------------+-------------------+------------+
    |age|  name|    app_list|      category_list|random_label|
    +---+------+------------+-------------------+------------+
    |  7| Alice| app_1,app_3|  cate_0:3,cate_1:2|           0|
    |  7|   Bob|app_3,app_21| cate_0:4,cate_10:7|           1|
    |  3|Claire| app_1,app_7| cate_0:3,cate_11:4|           0|
    |  7|   Dan| app_9,app_5|cate_20:2,cate_99:7|           0|
    +---+------+------------+-------------------+------------+
    <BLANKLINE>
    >>> calc = FeatureScoreCalculator(labelCol='random_label', labelRange=[0,1],
    ...     compoCols=['app_list', 'category_list'],
    ...     compoExploders=[explode_list_plain, explode_map_plain],
    ...     simpleCols=['age'])
    >>> result = calc.transform(df)
    >>> result.printSchema()
    root
     |-- feature: string (nullable = true)
     |-- info_value: double (nullable = true)
     |-- info_gain: double (nullable = true)
     |-- chi_sq: double (nullable = true)
    <BLANKLINE>
    """

    def __init__(self, labelCol: str, labelRange: List[int], compoCols: List[str], compoExploders: List[Callable],
                 simpleCols: List[str]):
        super().__init__()
        self.setLabelCol(labelCol) \
            .setLabelRange(labelRange) \
            .setCompoCols(compoCols) \
            .setCompoExploders(compoExploders) \
            .setSimpleCols(simpleCols)

    def setLabelCol(self, labelCol):
        """Sets the name of label column

        Args:
            labelCol: the name of label column

        Returns:
            ``self``
        """
        self._labelCol = labelCol
        return self

    def getLabelCol(self):
        """

        Returns:
            the name of label column
        """
        return self._labelCol

    def setLabelRange(self, labelRange):
        """Sets the range of label value.

        Args:
            labelRange: the range of label value

        Returns:
            ``self``
        """
        self._labelRange = labelRange
        return self

    def getLabelRange(self):
        """

        Returns:
            the range of label value
        """
        return self._labelRange

    def setCompoCols(self, compoCols):
        """Sets the names of composite columns.

        Args:
            compoCols: the names of composite columns

        Returns:
            ``self``
        """
        self._compoCols = compoCols
        return self

    def getCompoCols(self):
        """

        Returns:
            the names of composite columns
        """
        return self._compoCols

    def setCompoExploders(self, compoExploders):
        """Sets the content exploders of composite columns.

        Args:
            compoExploders: the content exploders of composite columns, in an order
                corresponding to the return value of ``getCompoCols()``

        Returns:
            ``self``
        """
        self._compoExploders = compoExploders
        return self

    def getCompoExploders(self):
        """

        Returns:
            the content exploders of composite columns
        """
        return self._compoExploders

    def setSimpleCols(self, simpleCols):
        """Sets the names of simple columns.

        Args:
            simpleCols: the names of simple columns

        Returns:
            ``self``
        """
        self._simpleCols = simpleCols
        return self

    def getSimpleCols(self):
        """

        Returns:
            the names of simple columns
        """
        return self._simpleCols

    def getInputCols(self) -> List[str]:
        """

        Returns:
            the names of input columns
        """
        return [self.getLabelCol()] + self.getCompoCols() + self.getSimpleCols()

    def getOutputCols(self) -> List[str]:
        """

        Returns:
            the names of newly generated columns
        """
        return [self.getLabelCol()] + ['feature'] + self.getAvailableScoreKinds()

    def isBinaryClass(self):
        """

        Returns:
            boolean indicator whether the label range is binary, that is,
            equal to ``{0, 1}`` or ``{-1, 1}``
        """
        if len(self.getLabelRange()) == 2:
            r = set(self.getLabelRange())
            return r == {0, 1} or r == {-1, 1}
        else:
            return False

    def getAvailableScoreKinds(self) -> List[str]:
        """

        Returns:
            the available kinds of scores under the constraint of ``labelRange``

        Notes:
            For multi-class records with more than 2 classes, information value
            is not available.
        """
        if self.isBinaryClass():
            return ['info_value', 'info_gain', 'chi_sq']
        else:
            return ['info_gain', 'chi_sq']

    def _transform(self, dataset):
        if self.isBinaryClass():
            result = _compute_scores_for_categorical_variable_binary_class(dataset,
                                                                           self.getLabelCol(),
                                                                           self.getLabelRange(),
                                                                           self.getCompoCols(),
                                                                           self.getCompoExploders(),
                                                                           self.getSimpleCols())
            assert set(self.getAvailableScoreKinds()).issubset(result.columns)
            return result
        else:
            result = _compute_scores_for_categorical_variable_multiclass(dataset,
                                                                         self.getLabelCol(),
                                                                         self.getLabelRange(),
                                                                         self.getCompoCols(),
                                                                         self.getCompoExploders(),
                                                                         self.getSimpleCols())
            assert set(self.getAvailableScoreKinds()).issubset(result.columns)
            return result


def _rank_with_column(df: DataFrame, input_col: str, output_col: str, ascending=True):
    column = df[input_col].asc() if ascending else df[input_col].desc()
    window = Window.orderBy(column)
    df = df.withColumn(output_col, functions.rank().over(window))
    return df


def _rank_with_columns(df: DataFrame, input_cols: List[str], output_cols: List[str], ascending=List[bool]) -> DataFrame:
    assert len(input_cols) == len(output_cols) and len(input_cols) == len(ascending)
    for i, o, a in zip(input_cols, output_cols, ascending):
        df = _rank_with_column(df, i, o, a)
    return df


class RowRanker(SparkTransformer, HasInputCols, HasOutputCols):
    """``RowRanker`` is a transformer that computes each row's rank in terms of
    each of the specified input columns. Note that the rank is computed independently
    in terms of each single input column.

    Args:
        inputCols: the list of the input column names
        outputCols: the list of the output column names, should match ``inputCols``
                in positions
        ascending: the list of options for ascending or descending orders, should
                match ``inputCols`` in positions

    >>> df = prepare_test_dataframe()
    >>> calc = FeatureScoreCalculator('random_label',
    ...                               [0,1],
    ...                               ['app_list', 'category_list'],
    ...                               [explode_list_plain, explode_map_plain],
    ...                               ['age'])
    >>> result = calc.transform(df)
    >>> ranker = RowRanker(inputCols=['info_value', 'info_gain', 'chi_sq'],
    ...                    outputCols=['info_value_rank', 'info_gain_rank', 'chi_sq_rank'],
    ...                    ascending=[False, False, False])
    >>> ranked = ranker.transform(result)
    >>> ranked.printSchema()
    root
     |-- feature: string (nullable = true)
     |-- info_value: double (nullable = true)
     |-- info_gain: double (nullable = true)
     |-- chi_sq: double (nullable = true)
     |-- info_value_rank: integer (nullable = true)
     |-- info_gain_rank: integer (nullable = true)
     |-- chi_sq_rank: integer (nullable = true)
    <BLANKLINE>
    >>> ranked.count()
    13

    """

    def __init__(self, inputCols: List[str], outputCols: List[str], ascending: List[bool]):
        super().__init__()
        self.setInputCols(inputCols)
        self.setOutputCols(outputCols)
        self.setAscending(ascending)

        assert len(self.getInputCols()) == len(self.getOutputCols())
        assert len(self.getInputCols()) == len(ascending)

    def setAscending(self, value: List[bool]):
        """Sets whether the scores are ranked in ascending order or the other way.

        Args:
            value: a list of booleans matching the order of ``inputCols``

        Returns:
            ``self``
        """
        self._ascending = value
        return self

    def getAscending(self):
        """

        Returns:
            a list of booleans matching the order of ``inputCols``, which was
            set by ``setAscending()``, indicating whether the ranks are ordered
            ascendingly or descendingly.
        """
        return self._ascending

    def _transform(self, dataset):
        return _rank_with_columns(dataset, self.getInputCols(), self.getOutputCols(), self.getAscending())


def _vectorize_row(row: Row, feature_dict: Dict[str, int], compo_cols: List[str], compo_exploders: List[Callable],
                   simple_cols: List[str]):
    pairs = []
    for name in simple_cols:
        if name in feature_dict:
            pairs.append((feature_dict[name], row[name]))
    for index, name in enumerate(compo_cols):
        explode = compo_exploders[index]
        exploded = explode(row[name])
        for key, value in exploded:
            if key in feature_dict:
                pairs.append((feature_dict[key], value))

    return Vectors.sparse(len(feature_dict), pairs)


def _vectorize_row_and_extend_old(dataset: DataFrame, feature_dict: Dict[str, int], label_col: str,
                                  compo_cols: List[str], compo_exploders: List[Callable], simple_cols: List[str],
                                  output_col):
    def vectorize(row: Row):
        return _vectorize_row(row, feature_dict, compo_cols, compo_exploders, simple_cols)

    def to_row(row: Row):
        d = row.asDict()
        d[output_col] = vectorize(row)
        return Row(**d)

    new_df = dataset.rdd.map(to_row).toDF()
    return new_df


class RowVectorizer(SparkTransformer, HasOutputCol):
    """``RowVectorizer`` is a ``Transformer`` that given a feature-index dictionary,
    composite column list and corresponding exploders, simple column list, and output
    column name, can transform a ``DataFrame`` to another with the specified columns
    assembled into a ``SparseVector``.


    Args:
        featureDict: a feature-index dictionary that maps a feature name to an index number
        labelCol: the name of label column
        compoCols: the list of composite columns you want to include
        compoExploders: a list of ``Callable``'s that each maps a composite cell to a list of ``(feature,value)`` tuples.
        simpleCols: the list of simple columns you want to include
        outputCol: the name of output column


    >>> df = prepare_test_dataframe()
    >>> df.printSchema()
    root
     |-- age: long (nullable = true)
     |-- name: string (nullable = true)
     |-- app_list: string (nullable = true)
     |-- category_list: string (nullable = true)
     |-- random_label: long (nullable = true)
    <BLANKLINE>
    >>> selected_features = ['age', 'app_1', 'cate_0']
    >>> feature_dict = {value:index for index, value in enumerate(selected_features)}
    >>> feature_dict
    {'age': 0, 'app_1': 1, 'cate_0': 2}
    >>> row_vectorizer = RowVectorizer(feature_dict,
    ...                                'random_label',
    ...                                ['app_list', 'category_list'],
    ...                                [explode_list_plain, explode_map_plain],
    ...                                ['age'],
    ...                                'features')
    >>> new_df = row_vectorizer.transform(df)
    >>> new_df.printSchema()
    root
     |-- age: long (nullable = true)
     |-- app_list: string (nullable = true)
     |-- category_list: string (nullable = true)
     |-- features: vector (nullable = true)
     |-- name: string (nullable = true)
     |-- random_label: long (nullable = true)
    <BLANKLINE>
    >>> new_df.count()
    4
    """

    def __init__(self, featureDict: Dict[str, int], labelCol: str, compoCols: List[str], compoExploders: List[Callable],
                 simpleCols: List[str], outputCol: str):
        super(RowVectorizer, self).__init__()
        self.setFeatureDict(featureDict) \
            .setLabelCol(labelCol) \
            .setCompoCols(compoCols) \
            .setCompoExploders(compoExploders) \
            .setSimpleCols(simpleCols) \
            .setOutputCol(outputCol)

    def setFeatureDict(self, featureDict):
        """

        Args:
            featureDict: the dictionary of features, mapping from the name of a
                feature to its index

        Returns:
            ``self``
        """
        self._featureDict = featureDict
        return self

    def getFeatureDict(self):
        """

        Returns:
            the feature dictionary
        """
        return self._featureDict

    def setLabelCol(self, labelCol):
        """Sets the name of label column

        Args:
            labelCol: the name of label column

        Returns:
            ``self``
        """
        self._labelCol = labelCol
        return self

    def getLabelCol(self):
        """

        Returns:
            the name of label column
        """
        return self._labelCol

    def setCompoCols(self, compoCols):
        """Sets the names of composite columns.

        Args:
            compoCols: the names of composite columns

        Returns:
            ``self``
        """
        self._compoCols = compoCols
        return self

    def getCompoCols(self):
        """

        Returns:
            the names of composite columns
        """
        return self._compoCols

    def setCompoExploders(self, compoExploders):
        """Sets the content exploders of composite columns.

        Args:
            compoExploders: the content exploders of composite columns, in an order
                corresponding to the return value of ``getCompoCols()``

        Returns:
            ``self``
        """
        self._compoExploders = compoExploders
        return self

    def getCompoExploders(self):
        """

        Returns:
            the content exploders of composite columns
        """
        return self._compoExploders

    def setSimpleCols(self, simpleCols):
        """Sets the names of simple columns.

        Args:
            simpleCols: the names of simple columns

        Returns:
            ``self``
        """
        self._simpleCols = simpleCols
        return self

    def getSimpleCols(self):
        """

        Returns:
            the names of simple columns
        """
        return self._simpleCols

    def getInputCols(self):
        """

        Returns:
            the names of input columns
        """
        return [self.getLabelCol()] + self.getCompoCols() + self.getSimpleCols()

    def getOutputCols(self):
        """

        Returns:
            the names of newly generated columns
        """
        return [self.getOutputCol()]

    @property
    def _column_appending(self):
        return True

    def _transform(self, dataset: DataFrame):
        return _vectorize_row_and_extend_old(dataset,
                                             self.getFeatureDict(),
                                             self.getLabelCol(),
                                             self.getCompoCols(),
                                             self.getCompoExploders(),
                                             self.getSimpleCols(),
                                             self.getOutputCol())


# Considerations:
#   0. Categorical variable of symbolic nature must be onehot encoded.
#      Categorical variable of numeric nature may be represented as numeric.
#   1. The conversion from categorical variable to onehot encoding needs special care.
#      Example: [('High',), ('Low',), ('Medium',)] ==> [(1,0), (0,0), (0,1)].
#      Consider using the combination of ``StringIndexer`` and ``OneHotEncoderEstimator``.
#   2. ``OneHotEncoderEstimator`` can produce a ``Transformer`` that transforms a column
#      of category indices to a column of binary vectors.
#   3. ``VectorAssembler`` supports generating ``Vector`` (sparse or dense, depending on
#      memory efficiency) from several columns.
#   4. Onehot encoding of categorical variable complicates some tasks following the
#      task of feature selection. One problem is that we may need to interpret the model
#      after we have the weight of each feature. In this case we need a mapping from
#      the index of the ``features`` vector to the name of represented feature. Onehot
#      encoding transforms a categorical variable to a multi-dimensional vector. And
#      ``VectorAssembler`` puts that vector as a part of another vector. We should track
#      the mapping between the indices.

class VanillaFeatureSelector(Estimator):
    """``VanillaFeatureSelector`` selects features from candidates based on user-specified
    simple criteria. Its ``fit()`` method generates a ``RowVectorizer`` that can convert
    specified columns to ``Vector`` form, preserving the selected features only.

    Args:
        labelCol: the name of label column
        labelRange: the collection of values that a label can take
        compoCols: the list of names of composite columns. A composite column represents a collection of features.
        compoExploders: the list of callables that disassemble each cell of composite columns into lists. The order of
            exploders should match that of ``compoCols``.
        simpleCols: the list of names of simple columns. A simple column represents a single feature.
        featuresCol: the name of the output column for the features vector
        selectionOption: the selection strategy option, see explanations below
        cutOff: the cutoff value accompanying ``selectionOption``, which means score cutoff or
            rank cutoff depending on ``selectionOption``

    ``VanillaFeatureSelector`` selects features based on one of the following statistics
    of all specified features:

    1. information value (available for binary class only),

       See also https://www.listendata.com/2015/03/weight-of-evidence-woe-and-information.html

       For a certain feature value ``a``, if the total number of events with that value
       or the total number of non-events with that value is zero, we use the adjustedWOE
       (see the above webpage) instead of WOE for that value ``a``.

    2. information gain,

       See also https://en.wikipedia.org/wiki/Information_gain_in_decision_trees

    3. chi square statistic.

       See also https://www.investopedia.com/terms/c/chi-square-statistic.asp

    It provides two selection strategies:

    1. score-based, with three options: ``info_value``, ``info_gain``, and ``chi_sq``,
    2. rank-based, with three options: ``info_value_rank``, ``info_gain_rank``, and ``chi_sq_rank``.

    The user can specify any of the above strategies by passing one of the above options
    as ``selectionOption`` along with the companion ``cutOff``. The score-based strategy
    preserves the features with the specified scores better than ``cutOff``. The rank-based
    strategy preserves the features with the specified rank better than ``cutOff``.
    For score-based strategy, larger scores means better quality; while for rank-based
    strategy, higher rank mean better quality.

    Here is some demo code:

    >>> df = prepare_test_dataframe()
    >>> df.printSchema()
    root
     |-- age: long (nullable = true)
     |-- name: string (nullable = true)
     |-- app_list: string (nullable = true)
     |-- category_list: string (nullable = true)
     |-- random_label: long (nullable = true)
    <BLANKLINE>
    >>> feature_selector = VanillaFeatureSelector(labelCol='random_label',
    ...                                           labelRange=[0, 1],
    ...                                           compoCols=['app_list', 'category_list'],
    ...                                           compoExploders=[explode_list_plain, explode_map_plain],
    ...                                           simpleCols=['age'],
    ...                                           featuresCol='features',
    ...                                           selectionOption='info_value_rank',
    ...                                           cutOff=5)
    >>> model = feature_selector.fit(df)
    >>> new_df = model.transform(df)
    >>> new_df.printSchema()
    root
     |-- age: long (nullable = true)
     |-- app_list: string (nullable = true)
     |-- category_list: string (nullable = true)
     |-- features: vector (nullable = true)
     |-- name: string (nullable = true)
     |-- random_label: long (nullable = true)
    <BLANKLINE>
    """

    def __init__(self, labelCol: str, labelRange: List[int], compoCols: List[str], compoExploders: List[Callable],
                 simpleCols: List[str], featuresCol: str, selectionOption: str, cutOff: float):
        super(VanillaFeatureSelector, self).__init__()
        self.setLabelCol(labelCol) \
            .setLabelRange(labelRange) \
            .setCompoCols(compoCols) \
            .setCompoExploders(compoExploders) \
            .setSimpleCols(simpleCols) \
            .setFeaturesCol(featuresCol) \
            .setSelectionOption(selectionOption) \
            .setCutOff(cutOff)

    def setLabelCol(self, labelCol):
        """Sets the name of label column

        Args:
            labelCol: the name of label column

        Returns:
            ``self``
        """
        self._labelCol = labelCol
        return self

    def getLabelCol(self):
        """

        Returns:
            the name of label column
        """
        return self._labelCol

    def setLabelRange(self, labelRange):
        """Sets the range of label value.

        Args:
            labelRange: the range of label value

        Returns:
            ``self``
        """
        self._labelRange = labelRange
        return self

    def getLabelRange(self):
        """

        Returns:
            the range of label value
        """
        return self._labelRange

    def setCompoCols(self, compoCols):
        """Sets the names of composite columns.

        Args:
            compoCols: the names of composite columns

        Returns:
            ``self``
        """
        self._compoCols = compoCols
        return self

    def getCompoCols(self):
        """

        Returns:
            the names of composite columns
        """
        return self._compoCols

    def setCompoExploders(self, compoExploders):
        """Sets the content exploders of composite columns.

        Args:
            compoExploders: the content exploders of composite columns, in an order
                corresponding to the return value of ``getCompoCols()``

        Returns:
            ``self``
        """
        self._compoExploders = compoExploders
        return self

    def getCompoExploders(self):
        """

        Returns:
            the content exploders of composite columns
        """
        return self._compoExploders

    def setSimpleCols(self, simpleCols):
        """Sets the names of simple columns.

        Args:
            simpleCols: the names of simple columns

        Returns:
            ``self``
        """
        self._simpleCols = simpleCols
        return self

    def getSimpleCols(self):
        """

        Returns:
            the names of simple columns
        """
        return self._simpleCols

    def setFeaturesCol(self, col_name):
        """Sets the name of the output column for feature vector. 

        Args:
            col_name: the column name of the output vector encoding features

        Returns:
            ``self``
        """
        self._featuresCol = col_name
        return self

    def getFeaturesCol(self):
        """

        Returns:
            the column name of the output vector encoding features
        """
        return self._featuresCol

    def getInputCols(self):
        """

        Returns:
            the list of input columns
        """
        return [self.getLabelCol()] + self.getCompoCols() + self.getSimpleCols()

    def isBinaryClass(self):
        """

        Returns:
            boolean indicator whether the label range is binary, that is,
            equal to ``{0, 1}`` or ``{-1, 1}``
        """
        if len(self.getLabelRange()) == 2:
            r = set(self.getLabelRange())
            return r == {0, 1} or r == {-1, 1}
        else:
            return False

    def getAvailableScoreKinds(self) -> List[str]:
        """

        Returns:
            the available kinds of scores under the constraint of ``labelRange``

        Notes:
            For multi-class records with more than 2 classes, information value
            is not available.
        """
        if self.isBinaryClass():
            return ['info_value', 'info_gain', 'chi_sq']
        else:
            return ['info_gain', 'chi_sq']

    def getRankOptions(self):
        """

        Returns:
            the names of available rank options
        """
        return [x + '_rank' for x in self.getScoreOptions()]

    def getScoreOptions(self):
        """

        Returns:
            the names of available score options
        """
        return self.getAvailableScoreKinds()

    def getAvailableSelectionOptions(self):
        """

        Returns:
            the names of all available selection options, both rank options and
            score options
        """
        return self.getRankOptions() + self.getScoreOptions()

    def setSelectionOption(self, option_name):
        """Sets the selection option.

        Args:
            option_name: the name of the selection option. Candidates are
                ``'info_value'``, ``'info_gain'``, ``'chi_sq'``,
                ``'info_value_rank'``, ``'info_gain_rank'``, and ``'chi_sq_rank'``.
                Note that ``'infor_value'`` and ``'info_value_rank'`` are available
                only for binary labels.

        Returns:
            ``self``
        """
        assert option_name in self.getAvailableSelectionOptions()
        self._selectionOption = option_name
        return self

    def getSelectionOption(self):
        """

        Returns:
            the selection option
        """
        return self._selectionOption

    def setCutOff(self, value):
        """Sets the cutoff value.

        Args:
            value: the cutoff value

        Returns:
            ``self``
        """
        self._cutOff = value
        return self

    def getCutOff(self):
        """

        Returns:
            the cutoff value
        """
        return self._cutOff

    def _fit(self, dataset):
        calculator = FeatureScoreCalculator(self.getLabelCol(),
                                            self.getLabelRange(),
                                            self.getCompoCols(),
                                            self.getCompoExploders(),
                                            self.getSimpleCols())
        scores = calculator.transform(dataset)
        score_kinds = calculator.getAvailableScoreKinds()

        scores.cache()

        if self.getSelectionOption() in self.getScoreOptions():
            filtered = scores.filter(scores[self.getSelectionOption()] >= self.getCutOff())
        else:
            rank_cols = [x + '_rank' for x in score_kinds]
            ascending = [False for x in score_kinds]
            ranker = RowRanker(score_kinds, rank_cols, ascending=ascending)
            ranked = ranker.transform(scores)
            filtered = ranked.filter(ranked[self.getSelectionOption()] <= self.getCutOff())

        features = filtered.select('feature').collect()
        features = [row['feature'] for row in features]
        feature_dict = {value: index for index, value in enumerate(features)}

        return RowVectorizer(feature_dict,
                             self.getLabelCol(),
                             self.getCompoCols(),
                             self.getCompoExploders(),
                             self.getSimpleCols(),
                             self.getFeaturesCol())


def batch_use_example():
    df = prepare_test_dataframe()
    df.show()
    feature_selector = VanillaFeatureSelector(labelCol='random_label',
                                              labelRange=[0, 1],
                                              compoCols=['app_list', 'category_list'],
                                              compoExploders=[explode_list_plain, explode_map_plain],
                                              simpleCols=['age'],
                                              featuresCol='features',
                                              selectionOption='info_value_rank',
                                              cutOff=5)
    model = feature_selector.fit(df)
    pprint(model.getFeatureDict())
    new_df = model.transform(df)
    new_df.show()


def interactive_use_example():
    df = prepare_test_dataframe()
    csc = FeatureScoreCalculator('random_label', [0, 1], ['app_list', 'category_list'],
                                 [explode_list_plain, explode_map_plain],
                                 ['age'])

    score_df = csc.transform(df)
    score_df.explain()
    score_df.persist()
    score_df.show()

    row_ranker = RowRanker(inputCols=['info_value', 'info_gain', 'chi_sq'],
                           outputCols=['info_value_rank', 'info_gain_rank', 'chi_sq_rank'],
                           ascending=[False, False, False])
    ranked_df = row_ranker.transform(score_df)
    ranked_df.explain()
    ranked_df.show()

    thresholdFeature = 'info_value_rank'  # specified by the data scientist
    thresholdRank = 5  # specified by the data scientist

    features = ranked_df.filter(ranked_df[thresholdFeature] <= thresholdRank).select('feature').collect()
    features = [row['feature'] for row in features]
    featureDict = {value: index for index, value in enumerate(features)}
    print(featureDict)

    row_vectorizer = RowVectorizer(featureDict=featureDict, labelCol='random_label',
                                   compoCols=['app_list', 'category_list'],
                                   compoExploders=[explode_list_plain, explode_map_plain],
                                   simpleCols=['age'],
                                   outputCol='features')
    new_df = row_vectorizer.transform(df)
    new_df.explain()
    new_df.show()


def groupwise_woe_and_iv_example():
    df = prepare_test_dataframe()
    result = compute_group_woe_iv(df,
                                  label_col='random_label',
                                  label_range=[0, 1],
                                  composite_cols=['app_list',
                                                  'category_list'],
                                  exploders=[explode_list_plain,
                                             explode_map_plain],
                                  simple_cols=['age'])

    result.show()


if __name__ == "__main__":
    # interactive_use_example()
    # batch_use_example()
    groupwise_woe_and_iv_example()
